import os
import torch
from torchvision import transforms, datasets
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.autograd import Variable
import numpy as np
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms, utils
import torch.optim as optim
from PIL import Image
from dataset import MyDataSet
from CNNet import MyCNN
from train import accuracy, train
from torch.utils.tensorboard import SummaryWriter
from tqdm import tqdm
from test import test
from detect import detect


# yourselves config
basePath = "D:\\myInterestTest\\objectDetect\\data\\catVSdog" # basePath 
dogFolderPath = basePath + os.sep + "train" + os.sep + "dog" # dog train images path
catFolderPath = basePath + os.sep + "train" + os.sep + "cat" # cat train images path
testImgPath = basePath + os.sep + "test" # 验证图片数据集
testPath = basePath + os.sep + "test.csv" # 验证图片类别csv文件
model_cp = './model2.pth' # 模型保存路径
tensorboard_path = 'D:\\myInterestTest\\objectDetect\\tensorBoard' # tensorboard 文件夹路径
dogAct = 0 # 狗的类别数字
catAct = 1 # 猫的类别数字
EPOCH = 20 # 训练轮数
workers = 10  # PyTorch读取数据线程数量
batch_size = 16 # 训练所抓取的数据样本数量


'''
  data_transform: 数据集转换设置
  注意：注意：注意：
    若改动 ：
    transforms.Resize(84), # 图片尺寸重设置
    transforms.CenterCrop(84), # 中心剪裁
    MyCNN 中的模型参数也需要相应进行改变

    改动位置：
    改动位置：
    改动位置：
    找到相应代码：
    self.out1 = nn.Linear(改动, 120) 61行
    改为：
    print(f"输出=>{len(x[0])}   长/宽=>{len(x[0][0])}") 查看输出 卷积层的输出
    它的输出。
    格式： 输出 * 长 * 宽 （长宽相等）
    格式： 输出 * 长 * 宽 （长宽相等）
    格式： 输出 * 长 * 宽 （长宽相等）
'''


data_transform = transforms.Compose([
    transforms.Resize(84), # 图片尺寸重设置
    transforms.CenterCrop(84), # 中心剪裁
    transforms.ToTensor(),
    transforms.Normalize((0.1307,), (0.3081,)),
])

if __name__ == '__main__':
    # 使用哪一种则注释其他两种

    # train 模型训练
    train(dogFolderPath=dogFolderPath, catFolderPath=catFolderPath, data_transform=data_transform, dogAct=dogAct, catAct=catAct,
          batch_size=batch_size, workers=workers, EPOCH=EPOCH, tensorboard_path=tensorboard_path,
          model_cp=model_cp)

    # test 模型验证
    test(model_path=model_cp, image_path=testImgPath, file_path=testPath, data_transform=data_transform, dogAct=dogAct, catAct=catAct,
         batch_size=batch_size, workers=workers)

    # detect 使用模型进行分类
    result = detect(model_path=model_cp, image_path='需要检测的图片或图片文件夹',
                    data_transform=data_transform, dogAct=dogAct, catAct=catAct, batch_size=batch_size, workers=1)
    print(result)
